import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix

# 防止中文乱码
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

# 1. 获取两年内10支股票的日线行情数据
tickers = ['920799.BJ', '920489.BJ', '920128.BJ', '920118.BJ', '920111.BJ',
           '920682.BJ', '920445.BJ', '920167.BJ', '920116.BJ', '920819.BJ']
start_date = '20220101'
end_date = '20240101'

all_data = []
for ticker in tickers:
    df = pro.daily(ts_code=ticker, start_date=start_date, end_date=end_date)
    df = df.sort_values('trade_date')
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    all_data.append(df)

combined_data = pd.concat(all_data, ignore_index=True)
combined_data.set_index('trade_date', inplace=True)


# 2. 对数据进行打标与计算
def calculate_technical_indicators(data):
    df = data.copy()
    # 移动平均线
    df['SMA_5'] = df['close'].rolling(window=5).mean()
    df['SMA_20'] = df['close'].rolling(window=20).mean()

    # RSI
    def rsi(data, period=14):
        deltas = data.diff()
        up = deltas.clip(lower=0)
        down = -deltas.clip(upper=0)
        avg_up = up.rolling(window=period).mean()
        avg_down = down.rolling(window=period).mean()
        rs = avg_up / avg_down
        rsi = 100 - (100 / (1 + rs))
        return rsi

    df['RSI'] = rsi(df['close'])
    # 布林带
    df['BB_upper'] = df['close'].rolling(window=20).mean() + 2 * df['close'].rolling(window=20).std()
    df['BB_lower'] = df['close'].rolling(window=20).mean() - 2 * df['close'].rolling(window=20).std()
    # MACD
    df['MACD'] = df['close'].ewm(span=12, adjust=False).mean() - df['close'].ewm(span=26, adjust=False).mean()
    df['Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
    return df.dropna()


technical_data = calculate_technical_indicators(combined_data)

# 技术指标可视化
ticker = tickers[0] # 以920799.BJ为例
stock_df = technical_data[technical_data['ts_code'] == ticker]

plt.figure(figsize=(12, 6))
plt.plot(stock_df.index, stock_df['close'], label='收盘价')
plt.plot(stock_df.index, stock_df['SMA_5'], label='SMA_5')
plt.plot(stock_df.index, stock_df['SMA_20'], label='SMA_20')
plt.title(f'{ticker} 收盘价与移动平均线')
plt.xlabel('日期')
plt.ylabel('价格')
plt.legend()
plt.show()

# 相对强弱指标（RSI）可视化
ticker = tickers[0] # 以920799.BJ为例
stock_df = technical_data[technical_data['ts_code'] == ticker]

plt.figure(figsize=(12, 6))
plt.plot(stock_df.index, stock_df['RSI'], label='RSI')
plt.axhline(y=70, color='r', linestyle='--', label='超买线 (70)')
plt.axhline(y=30, color='g', linestyle='--', label='超卖线 (30)')
plt.title(f'{ticker} 相对强弱指标 (RSI)')
plt.xlabel('日期')
plt.ylabel('RSI 值')
plt.legend()
plt.show()

# 3. 对打标后的数据进行处理与分析
# 空值处理
technical_data = technical_data.dropna()


# 异常值处理
def remove_outliers(df):
    num_cols = df.select_dtypes(include=[np.number]).columns
    Q1 = df[num_cols].quantile(0.25)
    Q3 = df[num_cols].quantile(0.75)
    IQR = Q3 - Q1
    df = df[~((df[num_cols] < (Q1 - 1.5 * IQR)) | (df[num_cols] > (Q3 + 1.5 * IQR))).any(axis=1)]
    return df


technical_data = remove_outliers(technical_data)

# 归一化
scaler = MinMaxScaler()
num_cols = technical_data.select_dtypes(include=[np.number]).columns
normalized_data = technical_data.copy()
normalized_data[num_cols] = scaler.fit_transform(technical_data[num_cols])

# 主成分分析
pca = PCA(n_components=0.95)
principal_components = pca.fit_transform(normalized_data[num_cols])
principal_df = pd.DataFrame(data=principal_components)

# 相关性分析
correlation_matrix = normalized_data[num_cols].corr()
plt.figure(figsize=(10, 8))
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm')
plt.title('相关矩阵')
plt.show()

# 降维后的数据可视化
plt.scatter(principal_df[0], principal_df[1])
plt.xlabel('主成分 1')
plt.ylabel('主成分 2')
plt.title('PCA')
plt.show()


# 个股画像
def stock_profile(df, stock, n=6):
    selected_stock = df[df['ts_code'] == stock]
    if not selected_stock.empty:
        # 筛选出数值类型的列
        num_cols = selected_stock.select_dtypes(include=[np.number]).columns
        last_row = selected_stock.iloc[-1][num_cols]
        # 确保 last_row 只包含数值类型数据
        last_row = pd.to_numeric(last_row, errors='coerce').dropna()
        if not last_row.empty:
            top_attributes = last_row.nlargest(n)
            bottom_attributes = last_row.nsmallest(n)
            print(f"Top {n} attributes for {stock}:")
            print(top_attributes)
            print(f"Bottom {n} attributes for {stock}:")
            print(bottom_attributes)
        else:
            print(f"No valid numerical data found for {stock}")


# 数据均衡（下采样）
# 假设分类标签：若次日收盘价上涨则为1，否则为0
normalized_data['label'] = (normalized_data['close'].shift(-1) > normalized_data['close']).astype(int)
normalized_data = normalized_data.dropna()
X = normalized_data[num_cols]
y = normalized_data['label']

# 找出正负样本数量
class_0_count, class_1_count = y.value_counts()
min_count = min(class_0_count, class_1_count)

# 下采样
class_0_indices = y[y == 0].index
class_1_indices = y[y == 1].index

class_0_downsampled = np.random.choice(class_0_indices, min_count, replace=False)
class_1_downsampled = np.random.choice(class_1_indices, min_count, replace=False)

downsampled_indices = np.concatenate([class_0_downsampled, class_1_downsampled])
X_resampled = X.loc[downsampled_indices]
y_resampled = y.loc[downsampled_indices]

# 4. 使用一种机器学习方法分析建模与模型评价
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_resampled, y_resampled, test_size=0.2, random_state=42)

# 随机森林分类器
clf = RandomForestClassifier(n_estimators=100, random_state=42)
clf.fit(X_train, y_train)

# 预测
y_pred = clf.predict(X_test)

# 评价指标
accuracy = accuracy_score(y_test, y_pred)
print(f"Accuracy: {accuracy}")
print(classification_report(y_test, y_pred))

# 混淆矩阵可视化
cm = confusion_matrix(y_test, y_pred)
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
plt.xlabel('预测值')
plt.ylabel('真实值')
plt.title('混淆矩阵')
plt.show()

